package com.zjj.lbw.ai.customer;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

/**
 * 智能客服系统
 */
public class CustomerServiceAgent {
    public static void main(String[] args) {
        // 加载并解析文件
        Document document;
        try {
            Path documentPath = Paths.get(CustomerServiceAgent.class.getClassLoader()
                    .getResource("meituan-qa.txt").toURI());
            DocumentParser documentParser = new TextDocumentParser();
            document = FileSystemDocumentLoader.loadDocument(documentPath, documentParser);

            // 切分文件
            DocumentSplitter splitter = new CustomerServiceDocumentSplitter() ;
            List<TextSegment> segments = splitter.split(document);

            // 文本向量化
            OpenAiEmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                    .baseUrl("http://langchain4j.dev/demo/openai/v1")
                    .apiKey("demo")
                    .modelName("text-embedding-3-small")
                    .build();
            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

            // 向量存储，可以只对 “问题“ 进行向量化，然后存储，因为搜索时，是只根据问题搜索，这样能提升一定的准确度
            EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
            embeddingStore.addAll(embeddings, segments);

            // 查找向量
//            String query = "余额怎么提现";
//            TextSegment textSegment = TextSegment.textSegment(query);
//            Embedding content = embeddingModel.embed(textSegment).content();
//            List<EmbeddingMatch<TextSegment>> embeddingMatchList =
//                    embeddingStore.findRelevant(content, 3, 0.7);
            // 可以用 ContentRetriever组件 检索向量
            ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(embeddingStore)
                    .embeddingModel(embeddingModel)
                    .maxResults(3) // 最相似的3个结果
                    .minScore(0.7) // 只找相似度在0.8以上的内容
                    .build();
            Query query = new Query("余额怎么提现");
            List<Content> contentList = contentRetriever.retrieve(query);

            // 将向量查找的结果，组装为 提示词，让大模型处理
            DefaultContentInjector injector = new DefaultContentInjector();
            UserMessage userMessage = injector.inject(contentList, UserMessage.from("余额怎么提现"));
            ChatLanguageModel model = OpenAiChatModel.builder().modelName("gpt-4o-mini").apiKey("demo").build();
            Response<AiMessage> response = model.generate(userMessage);
            System.out.println(response.content().text());

        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }

    static class CustomerServiceDocumentSplitter implements DocumentSplitter {

        @Override
        public List<TextSegment> split(Document document) {

            List<TextSegment> segments = new ArrayList<>();

            String[] parts = split(document.text());
            for (String part : parts) {
                segments.add(TextSegment.from(part));
            }

            return segments;
        }

        public String[] split(String text) {
            return text.split("\\s*\\R\\s*\\R\\s*");
        }

    }
}
